import java.io.File
import java.security.SecureRandom
import javax.crypto.Cipher
import javax.crypto.KeyGenerator
import javax.crypto.spec.SecretKeySpec
import org.bouncycastle.crypto.params.KeyParameter
import java.lang.management.ManagementFactory
import org.bouncycastle.crypto.engines.ChaChaEngine
import org.bouncycastle.crypto.params.ParametersWithIV

// ChaCha20 加密函数
fun chacha20Encrypt(data: ByteArray, key: ByteArray, iv: ByteArray): ByteArray {
    val engine = ChaChaEngine(20) // 20轮ChaCha20
    val params = ParametersWithIV(KeyParameter(key), iv)
    engine.init(true, params)
    val out = ByteArray(data.size)
    engine.processBytes(data, 0, data.size, out, 0)
    return out
}

// ChaCha20 解密函数（加密和解密过程一样）
fun chacha20Decrypt(data: ByteArray, key: ByteArray, iv: ByteArray): ByteArray {
    val engine = ChaChaEngine(20)
    val params = ParametersWithIV(KeyParameter(key), iv)
    engine.init(false, params)
    val out = ByteArray(data.size)
    engine.processBytes(data, 0, data.size, out, 0)
    return out
}

fun main() {
    // 1. 创建文件夹
    val dir = File("encrypt_bench")
    if (!dir.exists()) dir.mkdirs()

    // 2. 生成随机字节流
    val data = ByteArray(1024 * 1024) // 1MB
    SecureRandom().nextBytes(data)

    // 3. 初始化AES
    val keyGen = KeyGenerator.getInstance("AES")
    keyGen.init(128)
    val aesKey = keyGen.generateKey()
    val aesCipher = Cipher.getInstance("AES/ECB/PKCS5Padding")

    // 5. 监控CPU
    val osBean = ManagementFactory.getOperatingSystemMXBean() as com.sun.management.OperatingSystemMXBean

    val startTime = System.currentTimeMillis()
    val endTime = startTime + 5 * 60 * 1000 // 5分钟
    var count = 0

    var totalTime = 0L
    repeat(100) {
        // AES加密
        val t1 = System.currentTimeMillis()
        aesCipher.init(Cipher.ENCRYPT_MODE, aesKey)
        val aesEncrypted = aesCipher.doFinal(data)
        val t2 = System.currentTimeMillis()
        totalTime += (t2 - t1)
    }
    println("AES 总时间 $totalTime ms 平均耗时 ${totalTime.toFloat() / 10000}")
    totalTime = 0L

    // 生成密钥和IV
    val chachaKey = ByteArray(32) // ChaCha20 需要32字节密钥
    val chachaIv = ByteArray(8)   // ChaCha20 需要8字节IV
    SecureRandom().nextBytes(chachaKey)
    SecureRandom().nextBytes(chachaIv)

    repeat(100) {
        // ChaCha20加密
        val t3 = System.currentTimeMillis()
        // 加密
        val chachaEncrypted = chacha20Encrypt(data, chachaKey, chachaIv)
        val t4 = System.currentTimeMillis()
        totalTime += (t4 - t3)
    }
    println("ChaCha20加密 总时间 $totalTime ms 平均耗时 ${totalTime.toFloat() / 10000}")

    println("加密测试完成")
}
